#include "sweSolver.hpp"
#include "util.hpp"

namespace hms {

/** @brief Calculates new values for the state variables from the storage field */
template<typename MeshType>
void SWESolver<MeshType>:: calculateNewStateImp(){

	auto calcStateBlockwise = [&](Index colBeg, Index nCols){
		// /* option 1, loop with condition predeclared and d at the back */
		// Eigen::Block<ArrayXXs, Dynamic, Dynamic, true>
		// 	dBlock { d().values().middleCols(colBeg, nCols) },
		// 	vBlock { v().values().middleCols(colBeg, nCols) },
		// 	qBlock { q().values().middleCols(colBeg, nCols) };

		// auto dCond = qBlock.topRows<1>() < minWaterDepth;
		// for ( int i=0; i<2; ++i ){
		// 	vBlock.row(i) = dCond.select( 0, qBlock.row(i+1) / dBlock );
		// 	qBlock.row(i+1) = dCond.select( 0, qBlock.row(i+1) );
		// }

		// dBlock = qBlock.topRows<1>();


		/* option 2, like option 1 but no loop
		 * (Lenny: fastest version on my machine) */
		Array1XsMap dBlock { &d().values()(0, colBeg), nCols };
		Array2XsMap vBlock { &v().values()(0, colBeg), 2, nCols };
		Array3XsMap qBlock { &q().values()(0, colBeg), 3, nCols };

		auto dCond = qBlock.topRows<1>() < minWaterDepth;

		vBlock = dCond.replicate<2,1>()
			.select(0, qBlock.bottomRows<2>().rowwise() / qBlock.topRows<1>() );
		qBlock.bottomRows<2>() = dCond.replicate<2,1>()
			.select(0, qBlock.bottomRows<2>() );

		dBlock = qBlock.topRows<1>();

		// /* option 3, set d first, then use d as condition */
		// Eigen::Block<ArrayXXs, Dynamic, Dynamic, true>
		// 	dBlock { d().values().middleCols(colBeg, nCols) },
		// 	vBlock { v().values().middleCols(colBeg, nCols) },
		// 	qBlock { q().values().middleCols(colBeg, nCols) };

		// dBlock = qBlock.topRows<1>();

		// for ( int i=0; i<2; ++i ){
		// 	vBlock.row(i) = (dBlock < minWaterDepth)
		// 		.select( 0, qBlock.row(i+1) / dBlock );
		// 	qBlock.row(i+1) = (dBlock < minWaterDepth)
		// 		.select( 0, qBlock.row(i+1) );
		// }

		// /* option 4, like option 3, but without the loop */
		// Eigen::Block<ArrayXXs, Dynamic, Dynamic, true>
		// 	dBlock { d().values().middleCols(colBeg, nCols) },
		// 	vBlock { v().values().middleCols(colBeg, nCols) },
		// 	qBlock { q().values().middleCols(colBeg, nCols) };

		// dBlock = qBlock.topRows<1>();

		// vBlock = (dBlock < minWaterDepth).replicate<2,1>()
		// 	.select(0, qBlock.bottomRows<2>() / dBlock.replicate<2,1>() );
		// qBlock.bottomRows<2>() = (dBlock < minWaterDepth).replicate<2,1>()
		// 	.select(0, qBlock.bottomRows<2>() );
	};

	if constexpr ( std::is_base_of_v<StructMeshBase<MeshType>, MeshType> ){
		static constexpr Index ng {1};
		Index
			nx = this->mesh().nCellsX(),
			ny = this->mesh().nCellsY(),
			cellsPerRow { 2*ng + nx },
			domainCellsBeg { ng*(cellsPerRow+1) };
		
		OMP_PRAGMA( parallel for )
		for ( int j = 0; j < ny; ++j ){
			Index colBeg { domainCellsBeg + j*cellsPerRow };
			calcStateBlockwise(colBeg, nx);
		}
	} else {
		OMP_PRAGMA(parallel)
		{

		auto [colBeg, nCols] = ompSegment( this->mesh().nDomainCells() );
		
		calcStateBlockwise(colBeg, nCols);

		}
	}

	/* naive parallel version. slower. */
	// auto calcState = [&](Index cell){
	// 	d().values()(0,cell) = q().values()(0,cell);
	// 	for ( int i=0; i<2; ++i ){
	// 		v().values()(i,cell) = d().values()(0,cell) < minWaterDepth ?
	// 			0 : q().values()(i+1,cell) / d().values()(0,cell);
	// 		q().values()(i+1,cell) = d().values()(0,cell) < minWaterDepth ?
	// 			0 : q().values()(i+1,cell);
	// 	}
	// };
	
	// this->mesh().forEachDomainCellIndex( calcState );
}

template class SWESolver<Mesh>;
template class SWESolver<UniMesh>;
template class SWESolver<RectMesh>;
template class SWESolver<StructMesh>;

} // end namespace hms
